## This tests the normal and gamma distributions and their conjugacy
## relations

import random, sys, math

from pyblog.misc.util import Counter

from pyblog import *
from testmain import testmodule, MyTestCase

class testmeanprec(MyTestCase):
  """
  testmeanprec computes the mean and precision of a bunch of numbers using
  the normal normal-mean and gamma normal-prec conjugacy relations
  """
  def runTest(self):
    @var_dist
    def mean():
      return Normal(0, 1e-6)
    
    @var_dist
    def prec():
      return Gamma(.001, .001)
    
    @var_dist
    def item(i):
      return Normal(mean(), prec())
    
    # generate some data
    true_mean = 5
    true_prec = 2
    rand = random.Random(123456)
    data1 = [rand.normalvariate(true_mean, pow(true_prec, -0.5)) \
             for _ in range(1000)]
    
    ## compute sample avg and std. dev
    sample = Counter(dict((x,1) for x in data1))
    sample.normalize()

    pos_mean, pos_prec = query(
      [mean(), prec()],                           # query variables
      [item(i) == x for i, x in enumerate(data1)], # the observation
      burnin=10, scans=100, trace=disptrace, stats=dispstats)
    if dispstats:
      print "sample", sample.mean(), pow(sample.std_dev(), -2)
      print "pos_mean", pos_mean.mean(), pos_mean.std_dev()
      print "pos_prec", pos_prec.mean(), pos_prec.std_dev()
    self.assertAbsEqual(pos_mean.mean(), sample.mean(), sample.std_dev(),
                        "posterior mean")
    self.assertAlmostEqual(pow(sample.std_dev(), -2), pos_prec.mean(), 1,
                           "posterior prec")

class testlinreg(MyTestCase):
  """
  testlinreg: linear regression
  test @var also test Normal-Mean and Normal-Prec with Uniform prior
  """
  def runTest(self):
    @var
    def mu(i): return a() + b() * x[i]
    
    @var_dist
    def y(i): return Normal(mu(i), tau())
    
    @var_dist
    def a(): return Uniform(-100, 100)
    
    @var_dist
    def b(): return Uniform(-100, 100)
    
    @var_dist
    def tau(): return Uniform(0, 100)
    
    rand = random.Random(123456)

    # generate data around some fixed values for a, b, and tau
    true_a = 25
    true_b = 3
    true_tau = .25
    x = [rand.uniform(-100, 100) for _ in range(100)]
    data2 = [rand.normalvariate(true_a + true_b * z, pow(true_tau, -.5))\
             for z in x]


    pos_a, pos_b, pos_tau = query([a(), b(), tau()],
                                   [y(i) == z for i,z in enumerate(data2)],
                                   burnin = 10, scans=100, trace=False,
                                   stats=dispstats)
    if dispstats:
      print "pos_a", pos_a.mean(), pos_a.std_dev()
      print "pos_b", pos_b.mean(), pos_b.std_dev()
      print "pos_tau", pos_tau.mean(), pos_tau.std_dev()

    # check that the values are within 10% of the expected values
    self.assertPctEqual(pos_a.mean(), true_a, 10, "posterior a")
    self.assertPctEqual(pos_b.mean(), true_b, 10, "posterior b")
    self.assertPctEqual(pos_tau.mean(), true_tau, 10, "posterior tau")

class testhlinreg(MyTestCase):
  """
  Test heirarchical linear regression. Example taken from the "Rats" example
  in BUGS documentation. See:
     http://mathstat.helsinki.fi/openbugs/Examples/Rats.html
  """
  def runTest(self):
    
    @var_dist
    def Y(i,j): return Normal( mu(i,j), tau_c() )
    
    @var
    def mu(i, j):
      return alpha(i) + beta(i) * (x[j] - xbar)
    
    @var_dist
    def alpha(i): return Normal( alpha_c(), alpha_tau() )
    
    @var_dist
    def beta(i): return Normal( beta_c(), beta_tau() )
    
    @var_dist
    def tau_c(): return Gamma(.001, .001)

    @var
    def sigma(): return 1.0 / math.sqrt(tau_c())
    
    @var_dist
    def alpha_c(): return Normal(0, 1e-6)
    
    @var_dist
    def alpha_tau(): return Gamma(.001, .001)

    @var_dist
    def beta_c(): return Normal(0, 1e-6)

    @var_dist
    def beta_tau(): return Gamma(.001, .001)

    @var
    def alpha0(): return alpha_c() - xbar * beta_c()
    
    x = [8.0, 15.0, 22.0, 29.0, 36.0]
    xbar = 22
    N = 30
    T = 5
    data = [[151, 199, 246, 283, 320],
            [145, 199, 249, 293, 354],
            [147, 214, 263, 312, 328],
            [155, 200, 237, 272, 297],
            [135, 188, 230, 280, 323],
            [159, 210, 252, 298, 331],
            [141, 189, 231, 275, 305],
            [159, 201, 248, 297, 338],
            [177, 236, 285, 350, 376],
            [134, 182, 220, 260, 296],
            [160, 208, 261, 313, 352],
            [143, 188, 220, 273, 314],
            [154, 200, 244, 289, 325],
            [171, 221, 270, 326, 358],
            [163, 216, 242, 281, 312],
            [160, 207, 248, 288, 324],
            [142, 187, 234, 280, 316],
            [156, 203, 243, 283, 317],
            [157, 212, 259, 307, 336],
            [152, 203, 246, 286, 321],
            [154, 205, 253, 298, 334],
            [139, 190, 225, 267, 302],
            [146, 191, 229, 272, 302],
            [157, 211, 250, 285, 323],
            [132, 185, 237, 286, 331],
            [160, 207, 257, 303, 345],
            [169, 216, 261, 295, 333],
            [157, 205, 248, 289, 316],
            [137, 180, 219, 258, 291],
            [153, 200, 244, 286, 324]]

    # construct the observation
    obs = [Y(i,j) == data[i][j] for i in range(N) for j in range(T)]

    # construct the initial values
    init = [alpha(i) == 250 for i in range(N)] \
           + [beta(i) == 6 for i in range(N)] \
           + [alpha_c() == 150, beta_c() == 10, tau_c() == 1, alpha_tau() ==1,
              beta_tau() == 1]
    
    
    pos_alpha0, pos_beta_c, pos_sigma = \
                query([alpha0(), beta_c(), sigma()], obs, init,
                      burnin=10, scans=100, trace=disptrace,
                      stats=dispstats)

    if dispstats:
      print "alpha0", pos_alpha0.mean(), pos_alpha0.std_dev()
      print "beta.c", pos_beta_c.mean(), pos_beta_c.std_dev()
      print "sigma", pos_sigma.mean(), pos_sigma.std_dev()

    # compare with the values reported in BUGS and require the
    # absolute difference to the std. deviation reported in BUGS
    self.assertAbsEqual(pos_alpha0.mean(), 106.6, 3.655, "alpha0")
    self.assertAbsEqual(pos_beta_c.mean(), 6.185, 0.1061, "beta.c")
    self.assertAbsEqual(pos_sigma.mean(), 6.086, 0.4606, "sigma")

    # now, test with missing data

    # first we'll create a set of keys of observations to be deleted
    delvar = [Y(i,T-1) for i in range(5,10)] \
             + [Y(i,j) for i in range(10,20) for j in range(T-2, T)] \
             + [Y(i,j) for i in range(20,25) for j in range(T-3, T)] \
             + [Y(i,j) for i in range(25,30) for j in range(T-4, T)]
    delset = set([y.name() for y in delvar])
    obs2 = [y for y in obs if y.name() not in delset]

    # now query the last four observations of rat 25 and beta_c
    pos_y25_1, pos_y25_2, pos_y25_3, pos_y25_4, pos_beta_c = \
               query([Y(25,1), Y(25,2), Y(25,3), Y(25,4), beta_c()], obs2,
                     init, burnin=10, scans=100, trace=disptrace,
                     stats = dispstats)

    if dispstats:
      print "Y(25,1)", pos_y25_1.mean(), pos_y25_1.std_dev()
      print "Y(25,2)", pos_y25_2.mean(), pos_y25_2.std_dev()
      print "Y(25,3)", pos_y25_3.mean(), pos_y25_3.std_dev()
      print "Y(25,4)", pos_y25_4.mean(), pos_y25_4.std_dev()
      print "beta.c", pos_beta_c.mean(), pos_beta_c.std_dev()

    # again, we compare with the reported mean in BUGS and the absolute
    # difference to be within the reported std. deviation in BUGS
    self.assertAbsEqual(pos_y25_1.mean(), 204.6, 8.689, "Y(25,1)")
    self.assertAbsEqual(pos_y25_2.mean(), 250.2, 10.21, "Y(25,2)")
    self.assertAbsEqual(pos_y25_3.mean(), 295.6, 12.5, "Y(25,3)")
    self.assertAbsEqual(pos_y25_4.mean(), 341.2, 15.29, "Y(25,4)")
    self.assertAbsEqual(pos_beta_c.mean(), 6.578, 0.1497, "2: beta.c")
  
if __name__ == "__main__":
  testmodule(__import__("testnormalgamma"))
  
